import argparse
import torch
import random
from math import cos, pi

def load_table():
    age_map = {
        1: "Under 18",
        18: "18-24",
        25: "25-34",
        35: "35-44",
        45: "45-49",
        50: "50-55",
        56: "56+"
    }

    occupation_map = {
        0: "other",
        1: "academic/educator",
        2: "artist",
        3: "clerical/admin",
        4: "college/grad student",
        5: "customer service",
        6: "doctor/health care",
        7: "executive/managerial",
        8: "farmer",
        9: "homemaker",
        10: "K-12 student",
        11: "lawyer",
        12: "programmer",
        13: "retired",
        14: "sales/marketing",
        15: "scientist",
        16: "self-employed",
        17: "technician/engineer",
        18: "tradesman/craftsman",
        19: "unemployed",
        20: "writer"
    }
    return age_map, occupation_map

def cat2num(cat, type):
    age_map, occupation_map=load_table()
    #  构建反向映射：从 类别名 -> 编号（从1开始）
    if type == "age":
        # 提取所有类别名，按 key 排序保证顺序一致
        age_list = [age_map[k] for k in sorted(age_map.keys())]
        if cat in age_list:
            return age_list.index(cat) + 1  # 从1开始
        else:
            raise ValueError(f"Unknown age category: {cat}")
    elif type == "occupation":
        # 提取所有 occupation 类别，按 key 排序
        occ_list = [occupation_map[k] for k in sorted(occupation_map.keys())]
        if cat in occ_list:
            return occ_list.index(cat) + 1  # 从1开始
        else:
            raise ValueError(f"Unknown occupation category: {cat}")
    elif type=="gender":
        return 1 if cat=='Male' else 2
    else:
        raise ValueError("type must be 'age' or 'occupation'")
    
def get_args():
    parser = argparse.ArgumentParser()

    # Train params
    parser.add_argument('--batch_size', default=128, type=int)
    parser.add_argument('--lr', default=0.001, type=float)
    parser.add_argument('--maxlen', default=101, type=int)
    parser.add_argument('--warmup_steps', default=1000, type=int, help="Linear warmup over warmup_steps")
    parser.add_argument('--save_path', default="/home/fengyuebo/mnt/workspace/user-draw/result/sasrec_feat_gemini.pth", type=str, help="model save path")
    parser.add_argument('--type', default="train and test", type=str, help="model save path")

    # Baseline Model construction
    parser.add_argument('--hidden_units', default=64, type=int)
    parser.add_argument('--num_blocks', default=4, type=int)
    parser.add_argument('--num_epochs', default=100, type=int)
    parser.add_argument('--num_heads', default=4, type=int)
    parser.add_argument('--dropout_rate', default=0.2, type=float)
    parser.add_argument('--l2_emb', default=0, type=float)
    parser.add_argument('--device', default='cuda', type=str)
    parser.add_argument('--state_dict_path', default=None, type=str)
    parser.add_argument('--norm_first', default=True, type=bool)
    parser.add_argument('--feat', default=1, type=int)

    args = parser.parse_args()
    args.device='cuda' if torch.cuda.is_available() else 'cpu'

    return args

def collate_fn(batch):
    """
    将批次数据整理成模型可以使用的张量格式
    
    Args:
        batch: 一个批次的数据，每个元素是 (seq, pos, neg) 的元组
        
    Returns:
        seqs: 输入序列张量 [batch_size, maxlen]
        pos_seqs: 正样本序列张量 [batch_size, maxlen]  
        neg_seqs: 负样本序列张量 [batch_size, maxlen]
    """
    # 分离序列、正样本和负样本
    # print(batch)
    seqs, pos_seqs, neg_seqs, user_feat, item_feat = zip(*batch)

    # 转换为张量
    seqs = torch.stack(seqs)  # [batch_size, maxlen]
    pos_seqs = torch.stack(pos_seqs) # [batch_size, maxlen]
    neg_seqs = torch.stack(neg_seqs) # [batch_size, maxlen]
    
    return seqs, pos_seqs, neg_seqs, user_feat, item_feat

def model_init(model):
    for name, param in model.named_parameters():
        try:
            torch.nn.init.xavier_normal_(param.data)
        except:
            pass # just ignore those failed init layers
    model.pos_emb.weight.data[0, :] = 0
    model.item_emb.item_emb.weight.data[0, :] = 0
    
    # 正确初始化用户嵌入层的padding_idx
    if hasattr(model.user_emb, 'user_emb'):
        model.user_emb.user_emb.weight.data[0, :] = 0
    if hasattr(model.user_emb, 'age_emb'):
        model.user_emb.age_emb.weight.data[0, :] = 0
    if hasattr(model.user_emb, 'occupation_emb'):
        model.user_emb.occupation_emb.weight.data[0, :] = 0
    if hasattr(model.user_emb, 'gender_emb'):
        model.user_emb.gender_emb.weight.data[0, :] = 0
    
    # 初始化新添加的用户特征融合组件
    if hasattr(model, 'user_gate'):
        for module in model.user_gate:
            if hasattr(module, 'weight'):
                torch.nn.init.xavier_normal_(module.weight)
                if hasattr(module, 'bias') and module.bias is not None:
                    torch.nn.init.zeros_(module.bias)
    
    if hasattr(model, 'user_proj'):
        torch.nn.init.xavier_normal_(model.user_proj.weight)
        if model.user_proj.bias is not None:
            torch.nn.init.zeros_(model.user_proj.bias)
    
    if hasattr(model, 'user_item_interaction'):
        for name, param in model.user_item_interaction.named_parameters():
            if 'weight' in name:
                torch.nn.init.xavier_normal_(param)
            elif 'bias' in name:
                torch.nn.init.zeros_(param)
    
    if hasattr(model, 'user_enhancement'):
        for module in model.user_enhancement:
            if hasattr(module, 'weight'):
                torch.nn.init.xavier_normal_(module.weight)
                if hasattr(module, 'bias') and module.bias is not None:
                    torch.nn.init.zeros_(module.bias)
    
    return model

def warmup_cosine(optimizer, current_epoch, max_epoch, lr_min=0, lr_max=0.1, warmup_epoch = 1000):
    if current_epoch < warmup_epoch:
        lr = lr_max * current_epoch / warmup_epoch
    else:
        lr = lr_min + (lr_max-lr_min)*(1 + cos(pi * (current_epoch - warmup_epoch) / (max_epoch - warmup_epoch))) / 2
    for param_group in optimizer.param_groups:
        param_group["lr"] = lr

def cal_hr(model, seq, pos, item_num, user_feat=None, item_feat=None, k=10):
    # 拼接候选物品序列; 
    items=[]
    all_items = set(range(1, item_num + 1))
    for idx, p in enumerate(pos):
        # pos_set = set(p.numpy().tolist())
        # neg_candidates = list(all_items - pos_set)
        # items.append([p[-1]]+neg_candidates)
        last_p=p[-1]
        item=[last_p]+[i for i in range(1,last_p)]+[i for i in range(last_p+1, item_num+1)]
        items.append(item)
    
    # 计算logit
    logits=model.predict(seq, items, user_feat, None)
    _, indices = torch.topk(logits, k=k)
    num = ((indices == 0).any(dim=-1)).sum()
    hr=num.item()/logits.shape[0]
    
    return hr

def drop_zero(seq):
    mask = (seq != 0).any(dim=1) 
    if ((seq == 0).all(dim=1)).any() :
        print(seq)
        print(seq[mask])
    return seq[mask]